{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "width = 6\n",
    "height = 3\n",
    "import matplotlib\n",
    "matplotlib.rcParams['figure.figsize'] = [width, height]\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "## utilities\n",
    "import os\n",
    "\n",
    "## deep learning module\n",
    "import mxnet as mx\n",
    "\n",
    "## data processing\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "## reporting\n",
    "import perf\n",
    "from scipy.stats import pearsonr, spearmanr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.4.0\n"
     ]
    }
   ],
   "source": [
    "print(mx.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## some hyperparameters we won't tune via command line inputs\n",
    "DATA_SEGMENTS    = { 'tr': 0.6, 'va': 0.2, 'tst': 0.2}\n",
    "THRESHOLD_EPOCHS = 2\n",
    "COR_THRESHOLD    =  0.005\n",
    "\n",
    "## temporal slicing\n",
    "WIN              = 24 ##* 7\n",
    "H                = 3\n",
    "\n",
    "## model details \n",
    "MODEL            = 'rnn_model'\n",
    "SZ_FILT          = 8\n",
    "N_FILT           = 10\n",
    "RNN_UNITS        = 10\n",
    "SEASONAL_PERIOD  = 24\n",
    "\n",
    "## training details\n",
    "GPU              = 0\n",
    "BATCH_N          = 1024\n",
    "LR               = 0.0001\n",
    "DROP             = 0.2\n",
    "N_EPOCHS         = 30\n",
    "\n",
    "## data details\n",
    "DATA_FILE        = 'electricity.diff.txt'\n",
    "SAVE_DIR         = \"resultsDir\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: look at the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "elec = pd.read_csv('electricity.diff.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Handy data structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## courtesy of https://www.saltycrane.com/blog/2007/11/python-circular-buffer/\n",
    "class RingBuffer:\n",
    "    def __init__(self, size):\n",
    "        self.data = [0 for i in range(size)]\n",
    "\n",
    "    def append(self, x):\n",
    "        self.data.pop(0)\n",
    "        self.data.append(x)\n",
    "\n",
    "    def get(self):\n",
    "        return self.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################################\n",
    "## DATA PREPARATION ##\n",
    "################################\n",
    "\n",
    "def prepared_data(data_file, win, h, model_name):\n",
    "    df = pd.read_csv(data_file, sep=',', header=0)\n",
    "    x  = df.iloc[:, :].values ## need to drop first column as that's an index not a value\n",
    "    x = (x - np.mean(x, axis = 0)) / (np.std(x, axis = 0)) ## normalize data\n",
    "    \n",
    "    if model_name == 'fc_model':\n",
    "        ## provide first and second step lookbacks in one flat input\n",
    "        X = np.hstack([x[1:-h], x[0:-(h+1)]])\n",
    "        Y = x[(h+1):]\n",
    "        return (X, Y)\n",
    "    else:    \n",
    "        # preallocate X and Y data arrays\n",
    "        # X shape = num examples * time win * num channels (NTC)\n",
    "        X = np.zeros((x.shape[0] - win - h, win, x.shape[1]))\n",
    "        # Y shape = num examples * num channels\n",
    "        Y = np.zeros((x.shape[0] - win - h, x.shape[1]))\n",
    "        \n",
    "        for i in range(win, x.shape[0] - h):\n",
    "            y_i = x[i + h - 1     , :] ## the target value is h steps ahead\n",
    "            x_i = x[(i - win) : i , :] ## the input data are the previous win steps\n",
    "            X[i-win] = x_i\n",
    "            Y[i-win] = y_i\n",
    "\n",
    "        return (X, Y)\n",
    "\n",
    "\n",
    "def prepare_iters(data_file, win, h, model, batch_n):\n",
    "    X, Y = prepared_data(data_file, win, h, model)\n",
    "\n",
    "    n_tr = int(Y.shape[0] * DATA_SEGMENTS['tr'])\n",
    "    n_va = int(Y.shape[0] * DATA_SEGMENTS['va'])\n",
    "\n",
    "    X_tr, X_valid, X_test = X[                      : n_tr], \\\n",
    "                               X[n_tr             : n_tr + n_va], \\\n",
    "                               X[n_tr + n_va : ]\n",
    "    Y_tr, Y_valid, Y_test = Y[                      : n_tr], \\\n",
    "                               Y[n_tr             : n_tr + n_va], \\\n",
    "                               Y[n_tr + n_va : ]\n",
    "    \n",
    "    iter_tr = mx.io.NDArrayIter(data       = X_tr,\n",
    "                                   label      = Y_tr,\n",
    "                                   batch_size = batch_n)\n",
    "    iter_val = mx.io.NDArrayIter(  data       = X_valid,\n",
    "                                   label      = Y_valid,\n",
    "                                   batch_size = batch_n)\n",
    "    iter_test = mx.io.NDArrayIter( data       = X_test,\n",
    "                                   label      = Y_test,\n",
    "                                   batch_size = batch_n)\n",
    "\n",
    "    return (iter_tr, iter_val, iter_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################\n",
    "## MODELS ##\n",
    "################\n",
    "\n",
    "def fc_model(iter_train, input_feature_shape, X, Y,\n",
    "             win, sz_filt, n_filter, drop, seasonal_period):\n",
    "    output = mx.sym.FullyConnected(data=X, num_hidden=20)\n",
    "    output = mx.sym.Activation(output, act_type = 'relu')\n",
    "    output = mx.sym.FullyConnected(data=output, num_hidden=10)\n",
    "    output = mx.sym.Activation(output, act_type = 'relu')\n",
    "    output = mx.sym.FullyConnected(data=output, num_hidden=321)\n",
    "    \n",
    "    loss_grad = mx.sym.LinearRegressionOutput(data=output, label=Y)\n",
    "    return (loss_grad,\n",
    "            [v.name for v in iter_train.provide_data],\n",
    "            [v.name for v in iter_train.provide_label])    \n",
    "    \n",
    "def cnn_model(iter_train, input_feature_shape, X, Y,\n",
    "              win, sz_filt, n_filter, drop, seasonal_period):\n",
    "    conv_input = mx.sym.reshape(data=X, shape=(0, 1, win, -1)) \n",
    "    ## Convolution expects 4d input (N x channel x height x width)\n",
    "    ## in our case channel = 1 (similar to a black and white image\n",
    "    ## height = time and width = channels slash electric locations\n",
    "    \n",
    "    cnn_output = mx.sym.Convolution(data=conv_input,\n",
    "                                    kernel=(sz_filt,\n",
    "                                            input_feature_shape[2]),\n",
    "                                    num_filter=n_filter)\n",
    "    cnn_output = mx.sym.Activation(data=cnn_output, act_type='relu')\n",
    "    cnn_output = mx.sym.reshape(mx.sym.transpose(data=cnn_output,\n",
    "                                                 axes=(0, 2, 1, 3)),\n",
    "                                shape=(0, 0, 0)) \n",
    "    cnn_output = mx.sym.Dropout(cnn_output, p=drop)\n",
    "        \n",
    "    output = mx.sym.FullyConnected(data=cnn_output,\n",
    "                                   num_hidden=input_feature_shape[2])\n",
    "    loss_grad = mx.sym.LinearRegressionOutput(data=output, label=Y)\n",
    "    return (loss_grad,\n",
    "            [v.name for v in iter_train.provide_data],\n",
    "            [v.name for v in iter_train.provide_label])    \n",
    "\n",
    "    \n",
    "def rnn_model(iter_train, input_feature_shape, X, Y,\n",
    "              win, sz_filt, n_filter, drop, seasonal_period):\n",
    "    rnn_cells = mx.rnn.SequentialRNNCell()\n",
    "    rnn_cells.add(mx.rnn.GRUCell(num_hidden=RNN_UNITS))\n",
    "    rnn_cells.add(mx.rnn.DropoutCell(drop))\n",
    "    outputs, _ = rnn_cells.unroll(length=win, inputs=X, merge_outputs=False)\n",
    "    rnn_output = outputs[-1] # only take value from final unrolled cell for use later\n",
    "    \n",
    "    output = mx.sym.FullyConnected(data=rnn_output, num_hidden=input_feature_shape[2])\n",
    "    loss_grad = mx.sym.LinearRegressionOutput(data=output, label=Y)\n",
    "    return (loss_grad,\n",
    "            [v.name for v in iter_train.provide_data],\n",
    "            [v.name for v in iter_train.provide_label])    \n",
    "\n",
    "## simplifications to\n",
    "## https://github.com/apache/incubator-mxnet/blob/master/example/multivariate_time_series/src/lstnet.py\n",
    "def simple_lstnet_model(iter_train,  input_feature_shape, X, Y,\n",
    "                        win, sz_filt, n_filter, drop, seasonal_period):\n",
    "    ## must be 4d or 5d to use padding functionality\n",
    "    conv_input = mx.sym.reshape(data=X, shape=(0, 1, win, -1)) \n",
    "\n",
    "    ## convolutional element\n",
    "    ## we add padding at the end of the time win\n",
    "    cnn_output = mx.sym.pad(data=conv_input,\n",
    "                            mode=\"constant\",\n",
    "                            constant_value=0,\n",
    "                            pad_width=(0, 0,\n",
    "                                       0, 0,\n",
    "                                       0, sz_filt - 1, \n",
    "                                       0, 0))\n",
    "    cnn_output = mx.sym.Convolution(data=cnn_output,\n",
    "                                    kernel=(sz_filt,\n",
    "                                            input_feature_shape[2]),\n",
    "                                    num_filter=n_filter)\n",
    "    cnn_output = mx.sym.Activation(data=cnn_output, act_type='relu')\n",
    "    cnn_output = mx.sym.reshape(mx.sym.transpose(data=cnn_output,\n",
    "                                                 axes=(0, 2, 1, 3)),\n",
    "                                shape=(0, 0, 0))\n",
    "    cnn_output = mx.sym.Dropout(cnn_output, p=drop)\n",
    "\n",
    "    ## recurrent element\n",
    "    stacked_rnn_cells = mx.rnn.SequentialRNNCell()\n",
    "    stacked_rnn_cells.add(mx.rnn.GRUCell(num_hidden=RNN_UNITS))\n",
    "    outputs, _ = stacked_rnn_cells.unroll(length=win,\n",
    "                                          inputs=cnn_output,\n",
    "                                          merge_outputs=False)\n",
    "    rnn_output = outputs[-1] # only take value from final unrolled cell for use later\n",
    "    n_outputs = input_feature_shape[2]\n",
    "    cnn_rnn_model = mx.sym.FullyConnected(data=rnn_output,\n",
    "                                          num_hidden=n_outputs)\n",
    "\n",
    "    ## ar element\n",
    "    ar_outputs = []\n",
    "    for i in list(range(input_feature_shape[2])):\n",
    "        ar_series = mx.sym.slice_axis(data=X,\n",
    "                                      axis=2,\n",
    "                                      begin=i,\n",
    "                                      end=i+1)\n",
    "        fc_ar = mx.sym.FullyConnected(data=ar_series, num_hidden=1)\n",
    "        ar_outputs.append(fc_ar)\n",
    "    ar_model = mx.sym.concat(*ar_outputs, dim=1)\n",
    "\n",
    "    output = cnn_rnn_model + ar_model\n",
    "    loss_grad = mx.sym.LinearRegressionOutput(data=output, label=Y)\n",
    "    return (loss_grad,\n",
    "            [v.name for v in iter_train.provide_data],\n",
    "            [v.name for v in iter_train.provide_label])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################\n",
    "## TRAINING ##\n",
    "################\n",
    "\n",
    "def train(symbol, iter_train, valid_iter, iter_test,\n",
    "          data_names, label_names,\n",
    "          save_dir, gpu):\n",
    "    ## save training information/results \n",
    "    if not os.path.exists(SAVE_DIR):\n",
    "        os.makedirs(SAVE_DIR)\n",
    "    printFile = open(os.path.join(SAVE_DIR, 'log.txt'), 'w')\n",
    "    def print_to_file(msg):\n",
    "        print(msg)\n",
    "        print(msg, file = printFile, flush = True)\n",
    "    ## print_to_file(args) ## preserve configuation to enable hyperparameter optimization\n",
    "    ## archiving results header\n",
    "    print_to_file('Epoch     Training Cor     Validation Cor')\n",
    "\n",
    "\n",
    "    ## storing prior epoch's values to set an improvement threshold\n",
    "    ## terminates early if progress slow\n",
    "    buf     = RingBuffer(THRESHOLD_EPOCHS)\n",
    "    old_val = None\n",
    "\n",
    "    ## mxnet boilerplate\n",
    "    ## defaults to 1 gpu of which index is 0\n",
    "    ##devs = [mx.gpu(gpu)]\n",
    "    devs   = mx.cpu()\n",
    "    module = mx.mod.Module(symbol,\n",
    "                           data_names=data_names,\n",
    "                           label_names=label_names,\n",
    "                           context=devs)\n",
    "    module.bind(data_shapes=iter_train.provide_data,\n",
    "                label_shapes=iter_train.provide_label)\n",
    "    module.init_params(mx.initializer.Uniform(0.1))\n",
    "    module.init_optimizer(optimizer='adam',\n",
    "                          optimizer_params={'learning_rate': LR})\n",
    "\n",
    "    ## training\n",
    "    for epoch in range( N_EPOCHS):\n",
    "        iter_train.reset()\n",
    "        iter_val.reset()\n",
    "        for batch in iter_train:\n",
    "            module.forward(batch, is_train=True) # compute predictions\n",
    "            module.backward()                    # compute gradients\n",
    "            module.update()                      # update parameters\n",
    "\n",
    "        ## training results\n",
    "        train_pred  = module.predict(iter_train).asnumpy()\n",
    "        train_label = iter_train.label[0][1].asnumpy()\n",
    "        train_perf  = perf.write_eval(train_pred, train_label,\n",
    "                                      save_dir, 'train', epoch)\n",
    "\n",
    "        ## validation results\n",
    "        val_pred  = module.predict(iter_val).asnumpy()\n",
    "        val_label = iter_val.label[0][1].asnumpy()\n",
    "        val_perf = perf.write_eval(val_pred, val_label,\n",
    "                                   save_dir, 'valid', epoch)\n",
    "\n",
    "        print_to_file('%d         %f       %f ' % (epoch, train_perf['COR'], val_perf['COR']))\n",
    "        \n",
    "        if epoch > 0:                                # if we don't yet have measures of improvement, skip\n",
    "            buf.append(val_perf['COR'] - old_val) \n",
    "        if epoch > 2:                                # if we do have measures of improvement, check them\n",
    "            vals = buf.get()\n",
    "            # print(vals)\n",
    "            # print(COR_THRESHOLD)\n",
    "            vals = [v for v in vals if v != 0]\n",
    "            if sum([v < COR_THRESHOLD for v in vals]) == len(vals):\n",
    "                print_to_file('EARLY EXIT')\n",
    "                break\n",
    "        old_val = val_perf['COR']\n",
    "                \n",
    "    ## testing\n",
    "    test_pred  = module.predict(iter_test).asnumpy()\n",
    "    test_label = iter_test.label[0][1].asnumpy()\n",
    "    test_perf = perf.write_eval(test_pred, test_label, save_dir, 'tst', epoch)\n",
    "    print_to_file('\\n TESTING PERFORMANCE')\n",
    "    print_to_file(test_perf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data iterators\n",
    "iter_train, iter_val, iter_test = prepare_iters(DATA_FILE, WIN, H, MODEL, BATCH_N)    \n",
    "\n",
    "## prepare symbols\n",
    "input_feature_shape = iter_train.provide_data[0][1]    \n",
    "X                   = mx.sym.Variable(iter_train.provide_data[0].name)\n",
    "Y                   = mx.sym.Variable(iter_train.provide_label[0].name)\n",
    "    \n",
    "# set up model\n",
    "model_dict = {\n",
    "    'fc_model'            : fc_model,\n",
    "    'rnn_model'           : rnn_model,\n",
    "    'cnn_model'           : cnn_model,\n",
    "    'simple_lstnet_model' : simple_lstnet_model\n",
    "    }\n",
    "\n",
    "model = model_dict[MODEL]\n",
    "    \n",
    "symbol, data_names, label_names = model(iter_train,\n",
    "                                        input_feature_shape, X, Y,\n",
    "                                        WIN, SZ_FILT,\n",
    "                                        N_FILT, DROP, SEASONAL_PERIOD)\n",
    "\n",
    "## train \n",
    "train(symbol, iter_train, iter_val, iter_test, data_names, label_names, SAVE_DIR, GPU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: load the results and evaluate the model performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_true = pd.read_csv(\"resultsDir/valid_label_24.csv\", index_col=0)\n",
    "results_pred  = pd.read_csv(\"resultsDir/valid_pred_24.csv\", index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_true.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(results_true.iloc[:, 0], results_pred.iloc[:, 0])\n",
    "pearsonr(results_true.iloc[:, 0], results_pred.iloc[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(results_true.iloc[:, 25], results_pred.iloc[:, 25])\n",
    "print(pearsonr(results_true.iloc[:,25], results_pred.iloc[:, 25]))\n",
    "print(spearmanr(results_true.iloc[:,25], results_pred.iloc[:, 25]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(results_true.iloc[:, 50], results_pred.iloc[:, 50])\n",
    "print(pearsonr(results_true.iloc[:, 50], results_pred.iloc[:, 50]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(results_true.iloc[1800:2000, 50])\n",
    "plt.plot(results_pred.iloc[1800:2000, 50] * 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(results_true.iloc[1800:2000, 25])\n",
    "plt.plot(results_pred.iloc[1800:2000, 25] * 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(results_pred.iloc[1800:2000, 25])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: how does the model perform against the null model?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
